#pragma once
#include <dialects/Builtin/Builtin.h>
#include <dialects/Builtin/Dialect.h>
#include <dialects/Builtin/Prototype.h>
#include <dialects/Builtin/Types.h>
#include <dialects/PL/Dialect.h>
#include <dialects/PL/TypesAndProtos.h>
#include <lcir/ADT/ListNode.h>
#include <lcir/IR/ExprTemplate.h>
#include <lcir/IR/ExprTemps.h>
#include <lcir/IR/Expression.h>
#include <lcir/IR/LambdaProto.h>
#include <lcir/IR/Op.h>
#include <lcir/IR/Prototype.h>

#include <cstdint>
#include <functional>
#include <memory>
#include <ratio>
#include <type_traits>
#include <vector>
#include <numeric>

#include "dialects/Builtin/Types.h"
#include "lcir/IR/Constant.h"

namespace lcir::pl {
using namespace builtin;
static std::vector<Prototype *> deduceCondBlock(const std::vector<Prototype *> &operandProtos) {
  size_t lambdaIdx = operandProtos.size() - 1;
  size_t blockIdx = lambdaIdx - 1;
  auto *cond = operandProtos[blockIdx]->cast<CondBlockProto>();
  auto *lambda = operandProtos[lambdaIdx]->getValue()->cast<Lambda>();
  auto *ret = new CondBlockProto(lambda);
  setEdge(cond, ret);
  return {ret};
}

class IfTemp : public ExprTempCRTP<IfTemp> {
 public:
  IfTemp(Expression *expr) : ExprTempCRTP<IfTemp>(expr) {}
  Lambda *getLambda() { return operands[1]->getValue()->dyn_cast<Lambda>(); }
  Value *getCond() { return operands[0]->getValue(); }
};
/**
 * @brief 操作数0: bool, 操作数1: lambda
 */
class If : public OpCRTP<If, PlDialect, IfTemp, LiveNode> {
 public:
  If() : OpCRTP("If") {}
  bool checkType(const std::vector<Prototype *> &operandVals) override {
    if (operandVals.size() != 2) return false;
    // lhs是bool类型，rhs是lambda类型
    return operandVals[0]->getType()->isa<BoolType>() && operandVals[1]->isa<LambdaProto>();
  }

  std::vector<Prototype *> deduceRetType(const std::vector<Prototype *> &operandProtos) override {
    return {new CondBlockProto(operandProtos[1]->getValue()->cast<Lambda>())};
  }
};

/**
 *
 * @brief  bool lambda CondBlock
 */
class ElIf : public OpCRTP<ElIf, PlDialect, LiveNode> {
 public:
  ElIf() : OpCRTP("ElIf") {}
  bool checkType(const std::vector<Prototype *> &operandVals) override {
    if (operandVals.size() != 3) return false;
    // lhs是bool类型，rhs是lambda类型
    return operandVals[0]->getType()->isa<BoolType>() && operandVals[1]->isa<CondBlockProto>() &&
           operandVals[2]->isa<LambdaProto>();
  }
  std::vector<Prototype *> deduceRetType(const std::vector<Prototype *> &operandProtos) override {
    return deduceCondBlock(operandProtos);
  }
};

// @poerand lambda CondBlock
class ElTemp : public ExprTempCRTP<ElTemp> {
 public:
  ElTemp(Expression *expr) : ExprTempCRTP<ElTemp>(expr) {}
  Lambda *getLambda() { return operands[1]->getValue()->dyn_cast<Lambda>(); }
};
class El : public OpCRTP<El, PlDialect, ElTemp, LiveNode> {
 public:
  El() : OpCRTP("El") {}
  bool checkType(const std::vector<Prototype *> &operandVals) override {
    return operandVals.size() == 2 && operandVals[0]->isa<CondBlockProto>() && operandVals[1]->isa<LambdaProto>();
  }
  std::vector<Prototype *> deduceRetType(const std::vector<Prototype *> &operandProtos) override {
    return deduceCondBlock(operandProtos);
  }
};
/**
 * @brief 操作数0: bool, 操作数1: lambda
 */
class LoopTemp : public ExprTempCRTP<LoopTemp> {
 public:
  LoopTemp(Expression *expr) : ExprTempCRTP<LoopTemp>(expr) {}
  Lambda *getLambda() { return operands[1]->getValue()->dyn_cast<Lambda>(); }
  Value *getCond() { return operands[0]->getValue(); }
};

class Loop : public OpCRTP<Loop, PlDialect, LiveNode, LoopTemp> {
 public:
  Loop() : OpCRTP("Loop") {}
  bool checkType(const std::vector<Prototype *> &operandVals) override {
    if (operandVals.size() != 2) return false;
    // lhs是bool类型，rhs是lambda类型
    return operandVals[0]->getType()->isa<BoolType>() && operandVals[1]->isa<LambdaProto>();
  }
  std::vector<Prototype *> deduceRetType(const std::vector<Prototype *> &operandProtos) override { return {}; }
};


class Continue : public OpCRTP<Continue, PlDialect, LiveNode> {
 public:
  Continue() : OpCRTP("Continue") {}
  bool checkType(const std::vector<Prototype *> &operandVals) override { return true; }
  std::vector<Prototype *> deduceRetType(const std::vector<Prototype *> &operandProtos) override { return {}; }
};
class Break : public OpCRTP<Break, PlDialect, LiveNode> {
 public:
  Break() : OpCRTP("Break") {}
  bool checkType(const std::vector<Prototype *> &operandVals) override { return true; }
  std::vector<Prototype *> deduceRetType(const std::vector<Prototype *> &operandProtos) override { return {}; }
};


/**
 * @brief 第一个操作数是load的变量，剩下的是下标
 */
class ArrLoad : public OpCRTP<ArrLoad, PlDialect> {
 public:
  ArrLoad() : OpCRTP("ArrLoad") {}
  bool checkType(const std::vector<Prototype *> &operandVals) override { return operandVals.size() >= 1; }
  std::vector<Prototype *> deduceRetType(const std::vector<Prototype *> &operandProtos) override {
    Prototype *ret{nullptr};
    auto *type = operandProtos[0]->getType();
    if (type->isa<Int32Type>()) {
      ret = new builtin::Int32VarProto();
    } else if (type->isa<FloatType>()) {
      ret = new builtin::FloatVarProto();
    } else {
      ret = new builtin::UnkownVarProto();
    }
    return {ret};
  }
};

/**
 * @brief 赋值操作，操作数: 要赋的值， 目标变量，[索引。。。]
 */
class ArrStore : public OpCRTP<ArrStore, PlDialect> {
 public:
  ArrStore() : OpCRTP("ArrStore") {}
  bool checkType(const std::vector<Prototype *> &operandVals) override { return operandVals.size() > 2; }
  std::vector<Prototype *> deduceRetType(const std::vector<Prototype *> &operandProtos) override { return {}; }
};
/**
 * @brief 赋值操作，操作数: 左操作数是要赋的值， 右操作数是目标的变量，
 */
class Assign : public OpCRTP<Assign, PlDialect> {
 public:
  Assign() : OpCRTP("Assign") {}
  bool checkType(const std::vector<Prototype *> &operandVals) override { return operandVals.size() == 2; }
  std::vector<Prototype *> deduceRetType(const std::vector<Prototype *> &operandProtos) override { return {}; }
};

// 操作数为空
class DeclVar : public OpCRTP<DeclVar, PlDialect> {
  std::string varName{""};
  Type *type{nullptr};
  bool constVar{false};

 public:
  DeclVar(Type *type) : OpCRTP("Decl"), type(type) {}
  DeclVar(Type *type, std::string_view varName) : OpCRTP("Decl"), varName(varName), type(type) {}
  DeclVar(Type *type, std::string_view varName, bool constVar) : OpCRTP("Decl"), varName(varName), type(type), constVar(constVar) {}
  bool checkType(const std::vector<Prototype *> &operandVals) override { return operandVals.empty(); }

  std::vector<Prototype *> deduceRetType(const std::vector<Prototype *> &operandProtos) override {
    Prototype *ret{nullptr};
    if (type->isa<Int32Type>()) {
      ret = new builtin::Int32VarProto();
    } else if (type->isa<FloatType>()) {
      ret = new builtin::FloatVarProto();
    } else {
      ret = new builtin::UnkownVarProto();
    }
    ret->setInfo(varName);
    return {ret};
  }

  bool isConstVar() { return constVar; }
  Constant * fold(const std::vector<Constant *> &operandVals) override {
    if(!isConstVar())return nullptr;
    auto* result = getResult();
    auto userExpr = result->getSingleUser()->getExpr();
    if(!userExpr || !userExpr->isOp<pl::Assign>()) return nullptr;
    auto* c = userExpr->getOperandValue(0);
    return c->tryFold();
  }
};

class DeclArr : public OpCRTP<DeclArr, PlDialect> {
  Type *type{nullptr};
  std::string varName{""};
  std::vector<int64_t> shape;

 public:
  // DeclArr(Type *type) : OpCRTP("Decl"), type(type) {}
  DeclArr(Type *type, std::string_view varName, const std::vector<int64_t> shape)
      : OpCRTP("DeclArr"), type(type), varName(varName), shape(shape) {}
  bool checkType(const std::vector<Prototype *> &operandVals) override { return operandVals.empty(); }
  std::vector<int64_t> getShape() { return shape; }
  int64_t getElementNums() { return std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<int64_t>()); }
  std::vector<Prototype *> deduceRetType(const std::vector<Prototype *> &operandProtos) override {
    Prototype *ret{nullptr};
    if (type->isa<Int32Type>()) {
      ret = new builtin::Int32VarProto();
    } else if (type->isa<FloatType>()) {
      ret = new builtin::FloatVarProto();
    } else {
      ret = new builtin::UnkownVarProto();
    }
    std::stringstream ss;
    ss << varName << "[";
    for (auto [idx, s] : IndexRange(shape)) {
      ss << s;
      if (idx != shape.size() - 1) ss << ", ";
    }
    ss << "]";
    ret->setInfo(ss.str());
    return {ret};
  }
};

/*------------------------运算符--------------------------------------------*/
// 这个类类似模板，自动推导返回值类型（隐式转换）
// 目前允许浮点数和整数加法
static std::vector<lcir::Prototype *> deducArithRetType(const std::vector<lcir::Prototype *> &operandProtos) {
  // float相加，返回float
  auto *lhsTy = operandProtos[0]->getType();
  auto *rhsTy = operandProtos[1]->getType();
  Prototype *ret{nullptr};
  // int相加，返回int
  if (lhsTy->isa<builtin::FloatType>() && rhsTy->isa<builtin::FloatType>()) {
    ret = new Int32VarProto();
  } else if (lhsTy->isa<builtin::Int32Type>() && rhsTy->isa<builtin::Int32Type>()) {
    ret = new Int32VarProto();
  }
  // 如果是浮点数和整数相加，返回浮点类型
  if (lhsTy->isa<builtin::FloatType>() || rhsTy->isa<builtin::FloatType>()) {
    ret = new FloatVarProto();
  }
  return {ret};
}

static bool checkArithType(const std::vector<lcir::Prototype *> &operandVals) {
  if (operandVals.size() != 2) return false;
  for (auto *prop : operandVals) {
    if (!prop->getType()->isa<Int32Type>() && !prop->getType()->isa<FloatType>()) return false;
  }
  return true;
}

// 取反操作
class Neg : public OpCRTP<Neg, PlDialect> {
 public:
  Neg() : OpCRTP("Neg") {}
  bool checkType(const std::vector<lcir::Prototype *> &operandVals) override { return operandVals.size() == 1; }
  std::vector<lcir::Prototype *> deduceRetType(const std::vector<lcir::Prototype *> &operandProtos) override {
    auto *ty = operandProtos[0]->getType();
    Prototype *ret{nullptr};
    if (ty->isa<builtin::FloatType>()) {
      ret = new FloatVarProto();
    } else if (ty->isa<builtin::Int32Type>()) {
      ret = new Int32VarProto();
    }
    return {ret};
  }
};

// !操作
class Seqz : public OpCRTP<Seqz, PlDialect> {
 public:
  Seqz() : OpCRTP("Seqz") {}
  bool checkType(const std::vector<lcir::Prototype *> &operandVals) override { return operandVals.size() == 1; }
  std::vector<lcir::Prototype *> deduceRetType(const std::vector<lcir::Prototype *> &operandProtos) override {
    auto *ty = operandProtos[0]->getType();
    Prototype *ret{nullptr};
    if (ty->isa<builtin::FloatType>()) {
      ret = new FloatVarProto();
    } else if (ty->isa<builtin::Int32Type>()) {
      ret = new Int32VarProto();
    }
    return {ret};
  }
};

// 是0的话转为0，否则转为1（转换成bool变量）
class Inz : public OpCRTP<Inz, PlDialect> {
 public:
  Inz() : OpCRTP("Inz") {}
  bool checkType(const std::vector<lcir::Prototype *> &operandVals) override { return operandVals.size() == 1; }
  std::vector<lcir::Prototype *> deduceRetType(const std::vector<lcir::Prototype *> &operandProtos) override {
    return {new BoolVarProto()};
  };
};

enum ArithOpKind {
  AddOp = 0,
  SubOp = 1,
  MulOp = 2,
  DivOp = 3,
  ModOp = 4,
};

static constexpr const char *arithOpNames[] = {"Add", "Sub", "Mul", "Div", "Mod"};

template <template <class ArgTy> class BinaryOp>
static Constant *foldBinaryOp(Constant *lhs, Constant *rhs) {
  Constant *ret{nullptr};
  if (lhs->getProto<Int32ConstProto>() && rhs->getProto<Int32ConstProto>()) {
    auto lhsVal = lhs->getProto<Int32ConstProto>()->getNumber();
    auto rhsVal = rhs->getProto<Int32ConstProto>()->getNumber();
    BinaryOp<int> op;
    ret = Constant::get(std::make_unique<Int32ConstProto>(op(lhsVal, rhsVal)));
  } else if (lhs->getProto<FloatConstProto>() && rhs->getProto<FloatConstProto>()) {
    auto lhsVal = lhs->getProto<FloatConstProto>()->getNumber();
    auto rhsVal = rhs->getProto<FloatConstProto>()->getNumber();
    BinaryOp<float> op;
    ret = Constant::get(std::make_unique<FloatConstProto>(op(lhsVal, rhsVal)));
  } else if (lhs->getProto<BoolConstProto>() && rhs->getProto<BoolConstProto>()) {
    auto lhsVal = lhs->getProto<BoolConstProto>()->getNumber();
    auto rhsVal = rhs->getProto<BoolConstProto>()->getNumber();
    BinaryOp<bool> op;
    ret = Constant::get(std::make_unique<BoolConstProto>(op(lhsVal, rhsVal)));
  } else if (lhs->getProto<FloatConstProto>() && rhs->getProto<Int32ConstProto>()) {
    auto lhsVal = lhs->getProto<FloatConstProto>()->getNumber();
    auto rhsVal = rhs->getProto<Int32ConstProto>()->getNumber();
    BinaryOp<float> op;
    ret = Constant::get(std::make_unique<FloatConstProto>(op(lhsVal, rhsVal)));
  } else if (lhs->getProto<Int32ConstProto>() && rhs->getProto<FloatConstProto>()) {
    auto lhsVal = lhs->getProto<Int32ConstProto>()->getNumber();
    auto rhsVal = rhs->getProto<FloatConstProto>()->getNumber();
    BinaryOp<float> op;
    ret = Constant::get(std::make_unique<FloatConstProto>(op(lhsVal, rhsVal)));
  }

  return ret;
}
// 把求余数模板特化出来
static Constant *foldModeOp(Constant *lhs, Constant *rhs) {
  auto *lhsProto = lhs->getProto<Int32ConstProto>();
  auto *rhsProto = rhs->getProto<Int32ConstProto>();
  if (!lhsProto || !rhsProto) return nullptr;
  auto lhsVal = lhsProto->getNumber();
  auto rhsVal = rhsProto->getNumber();
  std::modulus<int> op;
  return Constant::get(std::make_unique<Int32ConstProto>(op(lhsVal, rhsVal)));
}

template <ArithOpKind OpKind>
class ArithOp : public OpCRTP<ArithOp<OpKind>, PlDialect, BinaryExpr> {
 public:
  ArithOp() : OpCRTP<ArithOp, PlDialect, BinaryExpr>(arithOpNames[OpKind]) {}

  bool checkType(const std::vector<lcir::Prototype *> &operandVals) override { return checkArithType(operandVals); }

  std::vector<lcir::Prototype *> deduceRetType(const std::vector<lcir::Prototype *> &operandProtos) override {
    return deducArithRetType(operandProtos);
  }
  Constant *fold(const std::vector<Constant *> &operands) override {
    auto *lhs = operands[0];
    auto *rhs = operands[1];
    if constexpr (OpKind == AddOp) {
      return foldBinaryOp<std::plus>(lhs, rhs);
    } else if constexpr (OpKind == SubOp) {
      return foldBinaryOp<std::minus>(lhs, rhs);
    } else if constexpr (OpKind == MulOp) {
      return foldBinaryOp<std::multiplies>(lhs, rhs);
    } else if constexpr (OpKind == DivOp) {
      return foldBinaryOp<std::divides>(lhs, rhs);
    } else if constexpr (OpKind == ModOp) {
      return foldModeOp(lhs, rhs);
    }
    return nullptr;
  }
};

using Add = ArithOp<AddOp>;
using Sub = ArithOp<SubOp>;
using Mul = ArithOp<MulOp>;
using Div = ArithOp<DivOp>;
using Mod = ArithOp<ModOp>;

enum CmpOpKind {
  GTOp = 0,
  GEOp = 1,
  LTOp = 2,
  LEOp = 3,
  EQOp = 4,
  UEOp = 5,
};

static constexpr const char *CmpOpNames[] = {"GT", "GE", "LT", "LE", "EQ", "UE"};

template <CmpOpKind OpKind>
class CmpOp : public OpCRTP<CmpOp<OpKind>, PlDialect, BinaryExpr> {
 public:
  CmpOp() : OpCRTP<CmpOp, PlDialect, BinaryExpr>(CmpOpNames[OpKind]) {}

  bool checkType(const std::vector<lcir::Prototype *> &operandVals) override { return checkArithType(operandVals); }

  std::vector<lcir::Prototype *> deduceRetType(const std::vector<lcir::Prototype *> &operandProtos) override {
    return {new BoolVarProto()};
  }
};
using GT = CmpOp<GTOp>;
using GE = CmpOp<GEOp>;
using LT = CmpOp<LTOp>;
using LE = CmpOp<LEOp>;
using EQ = CmpOp<EQOp>;
using UE = CmpOp<UEOp>;

enum LogicOpKind {
  AndOp = 0,
  OrOp = 1,
};

// @ 逻辑操作，包含和、或，操作数为布尔值
static constexpr const char *LogicOpNames[] = {"And", "Or"};
template <LogicOpKind OpKind>
class LogicOp : public OpCRTP<LogicOp<OpKind>, PlDialect, BinaryExpr> {
 public:
  LogicOp() : OpCRTP<LogicOp, PlDialect, BinaryExpr>(LogicOpNames[OpKind]) {}

  bool checkType(const std::vector<lcir::Prototype *> &operandVals) override {
    if (operandVals.size() != 2) return false;
    if (!operandVals[0]->getType()->isa<BoolType>() || !operandVals[1]->getType()->isa<BoolType>()) return false;
    return true;
  }

  std::vector<lcir::Prototype *> deduceRetType(const std::vector<lcir::Prototype *> &operandProtos) override {
    return {new BoolVarProto()};
  }
};

using And = LogicOp<AndOp>;
using Or = LogicOp<OrOp>;

// 变量操作
}  // namespace lcir::pl